import os

from openai import OpenAI
from typing import Optional, List, Dict, Union
from dotenv import load_dotenv

class LLM:
    """
    一个对 OpenAI 及其兼容接口进行封装的类，提供统一调用方式。

    支持 chat completions 和 embeddings 调用。
    """

    def __init__(
            self,
            api_key: Optional[str] = None,
            base_url: Optional[str] = None,
            model_name: Optional[str] = None
    ):
        """
        初始化 LLM 实例。

        :param api_key: Optional[str], API 密钥，必须通过环境变量或参数传递。
        :param base_url: Optional[str], 模型服务的基础 URL，默认从环境变量中读取。
        :param model_name: Optional[str], 使用的模型名称，默认从环境变量中读取。
        """
        self.api_key = api_key or os.getenv("OPENAI_API_KEY")
        if not self.api_key:
            raise ValueError("API Key 必须通过环境变量或参数提供")

        self.base_url = base_url or os.getenv("OPENAI_API_URL", "https://api-inference.modelscope.cn/v1/")
        self.model_name = model_name or os.getenv("OPENAI_MODEL_NAME", "deepseek-ai/DeepSeek-V3")

        self.client = OpenAI(
            api_key=self.api_key,
            base_url=self.base_url
        )

    def chat(
            self,
            messages: List[Dict[str, str]],
            temperature: float = 0.7,
            max_tokens: int = 1024,
            top_p: float = 1.0,
            stream: bool = False,
            **kwargs
    ) -> Union[str, object]:
        """
        调用聊天补全接口，适用于对话式生成任务。

        :param messages: list of dict, 包含对话历史的消息列表，格式为 [{"role": "user", "content": "..."}, ...]
        :param temperature: float, 控制输出随机性，值越大越随机，默认 0.7。
        :param max_tokens: int, 最大输出 token 数量，默认 1024。
        :param top_p: float, 核采样参数，控制 token 选择范围，默认 1.0。
        :param stream: bool, 是否启用流式输出，默认 False。
        :param kwargs: 其他可选参数，如 stop、presence_penalty 等。
        :return: str or object, 返回模型生成的文本或流式响应对象。
        """
        response = self.client.chat.completions.create(
            model=self.model_name,
            messages=messages,
            temperature=temperature,
            max_tokens=max_tokens,
            top_p=top_p,
            stream=stream,
            **kwargs
        )

        if stream:
            return response  # 返回 generator
        else:
            return response.choices[0].message.content

    def embedding(
            self,
            text: Union[str, List[str]],
            model: Optional[str] = None
    ) -> List[List[float]]:
        """
        获取文本的嵌入向量（Embedding），适用于语义检索、相似度计算等任务。

        :param text: str or list of str, 需要编码的文本内容。
        :param model: Optional[str], 使用的 Embedding 模型名称，默认使用初始化时指定的 model_name。
        :return: list of list of float, 文本对应的 embedding 向量。
        """
        input_texts = [text] if isinstance(text, str) else text
        response = self.client.embeddings.create(
            model=model if model else self.model_name,
            input=input_texts
        )
        return [item.embedding for item in response.data]

    def set_model(self, model_name: str):
        """
        动态切换当前使用的模型。

        :param model_name: str, 新的模型名称。
        """
        self.model_name = model_name

    def set_base_url(self, base_url: str):
        """
        动态更新模型服务的 base_url。

        :param base_url: str, 新的 base_url。
        """
        self.base_url = base_url
        self.client = OpenAI(api_key=self.api_key, base_url=self.base_url)
